#include "kernel_operator.h"

using namespace AscendC;

class KernelScatterElementsUint8
{
    using T = uint8_t;

public:
    __aicore__ inline KernelScatterElementsUint8() {}
    __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,
                                int32_t axis, uint8_t reduce,
                                int32_t size, int32_t var_dim, int32_t *var_ndarray,
                                int32_t upd_size,
                                int32_t ind_size, int32_t ind_dim, int32_t *ind_ndarray)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->size = size;
        this->ind_size = ind_size;
        this->upd_size = upd_size;
        this->ind_dim = ind_dim;
        this->var_dim = var_dim;
        this->axis = axis;
        this->reduce = reduce;

        this->var_ndarray = var_ndarray;
        this->ind_ndarray = ind_ndarray;

        varGm.SetGlobalBuffer((__gm__ T *)var, this->size);
        indGm.SetGlobalBuffer((__gm__ int32_t *)indices, this->ind_size);
        updGm.SetGlobalBuffer((__gm__ T *)updates, this->upd_size);
    }

    __aicore__ inline void Process()
    {
        if (this->reduce == 0)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                varGm.SetValue(new_idx, updGm.GetValue(i));
            }
        }
        else if (this->reduce == 1)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                int32_t temp_var = (int32_t)varGm.GetValue(new_idx);
                int32_t temp_upd = (int32_t)updGm.GetValue(i);
                varGm.SetValue(new_idx, (T)(temp_var + temp_upd));
            }
        }
        else if (this->reduce == 2)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                int32_t temp_var = (int32_t)varGm.GetValue(new_idx);
                int32_t temp_upd = (int32_t)updGm.GetValue(i);
                varGm.SetValue(new_idx, (T)(temp_var * temp_upd));
            }
        }
    }

private:
    __aicore__ inline int32_t *unravel_index(int32_t ind)
    {
        int32_t result[10];
        for (int i = this->ind_dim - 1; i >= 0; --i)
        {
            result[i] = ind % this->ind_ndarray[i];
            ind /= ind_ndarray[i];
        }
        return result;
    }

private:
    GlobalTensor<T> varGm;
    GlobalTensor<int32_t> indGm;
    GlobalTensor<T> updGm;

    int32_t size;
    int32_t var_dim;
    int32_t *var_ndarray;
    int32_t ind_size;
    int32_t upd_size;
    int32_t ind_dim;
    int32_t *ind_ndarray;
    int32_t axis;
    uint8_t reduce;
};

class KernelScatterElementsHalf
{
    using T = half;

public:
    __aicore__ inline KernelScatterElementsHalf() {}
    __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,
                                int32_t axis, uint8_t reduce,
                                int32_t size, int32_t var_dim, int32_t *var_ndarray,
                                int32_t upd_size,
                                int32_t ind_size, int32_t ind_dim, int32_t *ind_ndarray)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->size = size;
        this->var_dim = var_dim;
        this->ind_size = ind_size;
        this->upd_size = upd_size;
        this->ind_dim = ind_dim;
        this->axis = axis;
        this->reduce = reduce;

        this->var_ndarray = var_ndarray;
        this->ind_ndarray = ind_ndarray;

        varGm.SetGlobalBuffer((__gm__ T *)var, this->size);
        indGm.SetGlobalBuffer((__gm__ int32_t *)indices, this->ind_size);
        updGm.SetGlobalBuffer((__gm__ T *)updates, this->upd_size);
    }

    __aicore__ inline void Process()
    {
        if (this->reduce == 0)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                varGm.SetValue(new_idx, updGm.GetValue(i));
            }
        }
        else if (this->reduce == 1)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                float temp_var = (float)varGm.GetValue(new_idx);
                float temp_upd = (float)updGm.GetValue(i);
                varGm.SetValue(new_idx, (T)(temp_var + temp_upd));
            }
        }
        else if (this->reduce == 2)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                float temp_var = (float)varGm.GetValue(new_idx);
                float temp_upd = (float)updGm.GetValue(i);
                varGm.SetValue(new_idx, (T)(temp_var * temp_upd));
            }
        }
    }

private:
    __aicore__ inline int32_t *unravel_index(int32_t ind)
    {
        int32_t result[10];
        for (int i = this->ind_dim - 1; i >= 0; --i)
        {
            result[i] = ind % this->ind_ndarray[i];
            ind /= ind_ndarray[i];
        }
        return result;
    }

private:
    GlobalTensor<T> varGm;
    GlobalTensor<int32_t> indGm;
    GlobalTensor<T> updGm;

    int32_t size;
    int32_t var_dim;
    int32_t *var_ndarray;
    int32_t ind_size;
    int32_t upd_size;
    int32_t ind_dim;
    int32_t *ind_ndarray;
    int32_t axis;
    uint8_t reduce;
};

class KernelScatterElementsFloat
{
    using T = float;

public:
    __aicore__ inline KernelScatterElementsFloat() {}
    __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,
                                int32_t axis, uint8_t reduce,
                                int32_t size, int32_t var_dim, int32_t *var_ndarray,
                                int32_t upd_size,
                                int32_t ind_size, int32_t ind_dim, int32_t *ind_ndarray)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->size = size;
        this->var_dim = var_dim;
        this->ind_size = ind_size;
        this->upd_size = upd_size;
        this->ind_dim = ind_dim;
        this->axis = axis;
        this->reduce = reduce;

        this->var_ndarray = var_ndarray;
        this->ind_ndarray = ind_ndarray;

        varGm.SetGlobalBuffer((__gm__ T *)var, this->size);
        indGm.SetGlobalBuffer((__gm__ int32_t *)indices, this->ind_size);
        updGm.SetGlobalBuffer((__gm__ T *)updates, this->upd_size);
    }

    __aicore__ inline void Process()
    {
        if (this->reduce == 0)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                varGm.SetValue(new_idx, updGm.GetValue(i));
            }
        }
        else if (this->reduce == 1)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                T temp_var = varGm.GetValue(new_idx);
                T temp_upd = updGm.GetValue(i);
                varGm.SetValue(new_idx, temp_var + temp_upd);
            }
        }
        else if (this->reduce == 2)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                T temp_var = varGm.GetValue(new_idx);
                T temp_upd = updGm.GetValue(i);
                varGm.SetValue(new_idx, temp_var * temp_upd);
            }
        }
    }

private:
    __aicore__ inline int32_t *unravel_index(int32_t ind)
    {
        int32_t result[10];
        for (int i = this->ind_dim - 1; i >= 0; --i)
        {
            result[i] = ind % this->ind_ndarray[i];
            ind /= ind_ndarray[i];
        }
        return result;
    }

private:
    GlobalTensor<T> varGm;
    GlobalTensor<int32_t> indGm;
    GlobalTensor<T> updGm;

    int32_t size;
    int32_t var_dim;
    int32_t *var_ndarray;
    int32_t ind_size;
    int32_t upd_size;
    int32_t ind_dim;
    int32_t *ind_ndarray;
    int32_t axis;
    uint8_t reduce;
};

class KernelScatterElementsInt32
{
    using T = int32_t;

public:
    __aicore__ inline KernelScatterElementsInt32() {}
    __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,
                                int32_t axis, uint8_t reduce,
                                int32_t size, int32_t var_dim, int32_t *var_ndarray,
                                int32_t upd_size,
                                int32_t ind_size, int32_t ind_dim, int32_t *ind_ndarray)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->size = size;
        this->var_dim = var_dim;
        this->ind_size = ind_size;
        this->upd_size = upd_size;
        this->ind_dim = ind_dim;
        this->axis = axis;
        this->reduce = reduce;

        this->var_ndarray = var_ndarray;
        this->ind_ndarray = ind_ndarray;

        varGm.SetGlobalBuffer((__gm__ T *)var, this->size);
        indGm.SetGlobalBuffer((__gm__ int32_t *)indices, this->ind_size);
        updGm.SetGlobalBuffer((__gm__ T *)updates, this->upd_size);
    }

    __aicore__ inline void Process()
    {
        if (this->reduce == 0 || this->reduce == 2)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                varGm.SetValue(new_idx, updGm.GetValue(i));
            }
        }
        else if (this->reduce == 1)
        {
            for (int32_t i = 0; i < this->ind_size; i++)
            {
                int32_t *result = unravel_index(i);
                result[this->axis] = indGm.GetValue(i);
                int32_t new_idx = 0;
                for (int j = 0; j < this->var_dim; ++j)
                {
                    new_idx *= this->var_ndarray[j];
                    new_idx += result[j];
                }
                T temp_var = varGm.GetValue(new_idx);
                T temp_upd = updGm.GetValue(i);
                varGm.SetValue(new_idx, temp_var + temp_upd);
            }
        }
        // else if (this->reduce == 2)
        // {
        //     for (int32_t i = 0; i < this->ind_size; i++)
        //     {
        //         int32_t *result = unravel_index(i);
        //         result[this->axis] = indGm.GetValue(i);
        //         int32_t new_idx = 0;
        //         for (int j = 0; j < this->var_dim; ++j)
        //         {
        //             new_idx *= this->var_ndarray[j];
        //             new_idx += result[j];
        //         }
        //         T temp_var = varGm.GetValue(new_idx);
        //         T temp_upd = updGm.GetValue(i);
        //         varGm.SetValue(new_idx, temp_var * temp_upd);
        //     }
        // }
    }

private:
    __aicore__ inline int32_t *unravel_index(int32_t ind)
    {
        int32_t result[10];
        for (int i = this->ind_dim - 1; i >= 0; --i)
        {
            result[i] = ind % this->ind_ndarray[i];
            ind /= ind_ndarray[i];
        }
        return result;
    }

private:
    GlobalTensor<T> varGm;
    GlobalTensor<int32_t> indGm;
    GlobalTensor<T> updGm;

    int32_t size;
    int32_t var_dim;
    int32_t *var_ndarray;
    int32_t ind_size;
    int32_t upd_size;
    int32_t ind_dim;
    int32_t *ind_ndarray;
    int32_t axis;
    uint8_t reduce;
};

extern "C" __global__ __aicore__ void scatter_elements(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);

    if (TILING_KEY_IS(1))
    {
        KernelScatterElementsFloat op;
        op.Init(var, indices, updates,
                tiling_data.axis, tiling_data.reduce,
                tiling_data.size, tiling_data.var_dim, tiling_data.var_ndarray,
                tiling_data.upd_size,
                tiling_data.ind_size, tiling_data.ind_dim, tiling_data.ind_ndarray);
        op.Process();
    }
    else if (TILING_KEY_IS(2))
    {
        KernelScatterElementsInt32 op;
        op.Init(var, indices, updates,
                tiling_data.axis, tiling_data.reduce,
                tiling_data.size, tiling_data.var_dim, tiling_data.var_ndarray,
                tiling_data.upd_size,
                tiling_data.ind_size, tiling_data.ind_dim, tiling_data.ind_ndarray);
        op.Process();
    }
    else if (TILING_KEY_IS(3))
    {
        KernelScatterElementsHalf op;
        op.Init(var, indices, updates,
                tiling_data.axis, tiling_data.reduce,
                tiling_data.size, tiling_data.var_dim, tiling_data.var_ndarray,
                tiling_data.upd_size,
                tiling_data.ind_size, tiling_data.ind_dim, tiling_data.ind_ndarray);
        op.Process();
    }
    else if (TILING_KEY_IS(4))
    {
        KernelScatterElementsUint8 op;
        op.Init(var, indices, updates,
                tiling_data.axis, tiling_data.reduce,
                tiling_data.size, tiling_data.var_dim, tiling_data.var_ndarray,
                tiling_data.upd_size,
                tiling_data.ind_size, tiling_data.ind_dim, tiling_data.ind_ndarray);
        op.Process();
    }
}
